{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 简体中文 | [English](./QAOA_En.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 准备\n",
    "\n",
    "本文档演示 Paddle Quantum 上量子近似优化算法（QAOA，Quantum Approximate Optimization Algorithm）的工作流程 [1]。\n",
    "\n",
    "开始之前完成准备工作：\n",
    "\n",
    "   - 调用飞桨 paddlepaddle\n",
    "\n",
    "   - 调用常用的库, 例如画图工具库 networkx 和 matplotlib.pyplot\n",
    "\n",
    "   - 调用自定义函数 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from paddle import fluid\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "from numpy import matmul as np_matmul\n",
    "from paddle.complex import matmul as pp_matmul\n",
    "from paddle.complex import transpose\n",
    "from paddle_quantum.circuit import UAnsatz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 背景\n",
    "\n",
    "\n",
    "量子近似优化算法（QAOA，Quantum Approximate Optimization Algorithm）是可以在近期有噪中等规模（NISQ，Noisy Intermediate-Scale Quantum）量子计算机上运行且具有广泛应用前景的量子算法。例如，QAOA 可以用来处理压缩图信号和二次优化等领域常见的离散组合优化问题。这类优化问题通常可以归结为下面的数学模型：\n",
    "\n",
    "\n",
    " $$F=\\max_{z_i\\in\\{-1,1\\}} \\sum q_{ij}(1-z_iz_j)=-\\min_{z_i\\in\\{-1,1\\}} \\sum q_{ij}z_iz_j+ \\sum q_{ij}. $$\n",
    "\n",
    "\n",
    "其中, $z_i \\in \\{-1 ,1\\} $ 是待求的二元参数，系数 $q_{ij}$ 是 $z_i z_j$ 的权重 (weight)。一般地，精确求解该问题对于经典计算机是 NP-hard 的，而 QAOA 被认为对近似求解这类困难问题具有潜在速度优势。\n",
    "\n",
    "QAOA 的工作原理是把上述经典优化问题（例如组合优化问题）甚至量子优化问题（例如量子多体系统中 Ising 模型的求解）等价地转化为求解一个物理系统哈密顿量（Hamiltonian）的基态能量（对应优化问题的最优值）及其相应的基态（对应于优化问题的最优解）。在数学形式上，QAOA 等价于求解一个实对角矩阵 $H$ 的最小特征值及其对应的特征向量。\n",
    "\n",
    "和另外一种常见的变分量子特征求解器(VQE, Variational Quantum Eigensolver) 一样，QAOA 也是一种量子-经典混杂算法。 然而 QAOA 参数化量子电路的实现更简单，仅需两个可以调节参数的量子电路模块组成。\n",
    "\n",
    "接下来，我们通过图的最大割问题 （Max-Cut problem）来展示 QAOA 算法的工作流程和原理。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# 示例\n",
    "\n",
    "##  1. Max-Cut 问题\n",
    "\n",
    "图的 Max-Cut 问题可以描述为：对于一个给定的包含 $N$ 个顶点 （nodes or vertices）和 $M$ 条边 (edges) 的无向图，找到一种分割方案将图的顶点集合分割成两个无交子集合 $S$ 和 $S^\\prime$，使得连接这两个顶点集合之间边的数目尽可能多。如图所示，我们考虑含4个顶点且具有环形结构的图： \n",
    "\n",
    " ![ring4.png](https://release-data.cdn.bcebos.com/PIC%2FMaxCut.png) \n",
    "\n",
    "Max-Cut 问题建模：在做分割时，若顶点 $i$ 属于集合 $S$ 时，赋值 $z_i=1$；若顶点 $j$ 属于 $S^\\prime$ 时，则令 $z_j=-1$。那么对于图的任意连接顶点 $(i, j)$ 的边则满足：若顶点属于同一集合 $S$ 或 $S^\\prime$ 时，$z_iz_j=1$; 若顶点分别属于不同集合时，$z_izj=-1$。于是 Max-Cut 问题转化为如下优化问题：\n",
    "\n",
    "$$ F=\\min_{z_i\\in\\{-1, 1\\}} z_1 z_2+z_2z_3+z_3z_4+z_4z_1.$$\n",
    "\n",
    "这里所有 $q_{ij}$ 均设置为 1，表示每条边的权重相等。该问题的所有可行解由比特串 $ \\boldsymbol{z}=z_1z_2z_3z_4 \\in \\{-1, 1\\}^4$ 组成，而且通常需要遍历所有比特串才能得到问题的最小值（最优解）。容易看出，比特串的数目是顶点数目 $N$ 的指数级别，即 $2^N$。因此，随着顶点数目的增加，搜索的代价也会呈指数级别增加。\n",
    "\n",
    "接下来，我们提供两种方法来预处理编码经典优化问题的图，即如何通过 Paddle Quantum 输入和可视化无权（或带权重）图：\n",
    "\n",
    "- 方法1是通过指定图的顶点和相应的边（及其权重）\n",
    "- 方法2是通过直接输入图的邻接矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def generate_graph(N, GRAPHMETHOD):\n",
    "    \"\"\"\n",
    "    It plots an N-node graph which is specified by Method 1 or 2.\n",
    "    \n",
    "    Args:\n",
    "        N: number of nodes (vertices) in the graph\n",
    "        METHOD: choose which method to generate a graph\n",
    "    Return:\n",
    "        the specific graph and its adjacency matrix\n",
    "    \"\"\"\n",
    "    # Method 1 generates a graph by self-definition\n",
    "    if GRAPHMETHOD == 1:\n",
    "        print(\"Method 1 generates the graph from self-definition using EDGE description\")\n",
    "        graph = nx.Graph()\n",
    "        graph_nodelist=range(N)\n",
    "        graph.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 0)])\n",
    "        graph_adjacency = nx.to_numpy_matrix(graph, nodelist=graph_nodelist)\n",
    "    # Method 2 generates a graph by using its adjacency matrix directly\n",
    "    elif GRAPHMETHOD == 2:\n",
    "        print(\"Method 2 generates the graph from networks using adjacency matrix\")\n",
    "        graph_adjacency = np.array([[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0]])\n",
    "        graph = nx.Graph(graph_adjacency)\n",
    "    else:\n",
    "        print(\"Method doesn't exist \")\n",
    "\n",
    "    return graph, graph_adjacency"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里指定方法1来预处理图：\n",
    "\n",
    "- 图的顶点数目 $N=4$\n",
    "- 图的输入方法 GRAPHMETHOD = 1 \n",
    "\n",
    "注意：上述两种方法给出的图的顶点均从 $0$ 开始计数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# number of qubits or number of nodes in the graph\n",
    "N=4  \n",
    "classical_graph, classical_graph_adjacency= generate_graph(N, GRAPHMETHOD=1)\n",
    "print(classical_graph_adjacency)\n",
    "\n",
    "pos = nx.circular_layout(classical_graph)\n",
    "nx.draw(classical_graph, pos, width=4, with_labels=True, font_weight='bold')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 编码量子优化问题\n",
    "\n",
    "接下来需要把上述经典优化问题映射为量子优化问题。利用替代关系 $z=1\\rightarrow |0\\rangle = \\begin{bmatrix}1 \\\\ 0\\end{bmatrix}$ 和 $z=-1\\rightarrow |1\\rangle = \\begin{bmatrix}0 \\\\ 1\\end{bmatrix}$, 我们把二元参数 $z_i\\in\\{-1, 1\\}$ 对应为描述量子比特的 Pauli-Z 算符 $Z_i=\\begin{bmatrix} 1 & 0\\\\ 0 & -1\\end{bmatrix}$ 的两个本征值。于是经典优化问题里的目标函数相应地编码为一个描述系统哈密顿量的矩阵\n",
    "\n",
    "$$H_{c}= Z_1Z_2+Z_2Z_3+Z_3Z_4+Z_4Z_1.$$\n",
    "\n",
    "其中 $Z_iZ_{j}$ 是 tensor product 运算，表示 Pauli-Z 算符分别作用在量子比特 $i$ 和 $j$ 上，而其余的量子比特上作用单位算符 $I=\\begin{bmatrix} 1 & 0\\\\ 0 & 1\\end{bmatrix}$ ，例如 $Z_1Z_2 =Z_1\\otimes Z_2\\otimes I_3\\otimes I_4$。经过上述操作，我们把经典优化问题转化为求解矩阵 $H_{c}$ 的最小特征值 $F$ 及其对应的向量 $|\\psi\\rangle$, 即\n",
    "\n",
    "$$ F=\\min_{|\\psi\\rangle} \\langle \\psi| H_c |\\psi\\rangle.$$\n",
    "\n",
    "这里，$|\\psi\\rangle$ 记为一个模长为1的 $2^4=16$ 维复向量，$\\langle \\psi|$ 是其共轭转置。\n",
    "\n",
    "Paddle Quantum 中通过函数 H_generator 完成编码任务："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def H_generator(N, adjacency_matrix):\n",
    "    \"\"\"\n",
    "    This function maps the given graph via its adjacency matrix to the corresponding Hamiltiona H_c.\n",
    "    \n",
    "    Args:\n",
    "        N: number of qubits, or number of nodes in the graph, or number of parameters in the classical problem\n",
    "        adjacency_matrix:  the adjacency matrix generated from the graph encoding the classical problem\n",
    "    Return:\n",
    "        H_graph: the problem-based Hamiltonian H generated from the graph_adjacency matrix for the given graph\n",
    "        H_graph_diag: the real part of the problem-based Hamiltonian H_graph\n",
    "    \"\"\"\n",
    "\n",
    "    sigma_Z = np.array([[1, 0], [0, -1]])\n",
    "    H = np.zeros([2 ** N, 2 ** N])\n",
    "    # Generate the Hamiltonian H_c from the graph via its adjacency matrix\n",
    "    for row in range(N):\n",
    "        for col in range(N):\n",
    "            if abs(adjacency_matrix[N - row - 1, N - col - 1]) and row < col:\n",
    "                identity_1 = np.diag(np.ones([2 ** row]))\n",
    "                identity_2 = np.diag(np.ones([2 ** (col - row - 1)]))\n",
    "                identity_3 = np.diag(np.ones([2 ** (N - col - 1)]))\n",
    "                H += adjacency_matrix[N - row - 1, N - col - 1] * np.kron(\n",
    "                    np.kron(np.kron(np.kron(identity_1, sigma_Z), identity_2), sigma_Z),\n",
    "                    identity_3,\n",
    "                )\n",
    "\n",
    "    H_graph = H.astype(\"complex64\")\n",
    "    H_graph_diag = np.diag(H_graph).real\n",
    "\n",
    "    return H_graph, H_graph_diag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以查看生成矩阵 $H_c $ 的具体形式，并且获取它的特征值信息："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "_, H_problem_diag = H_generator(N, classical_graph_adjacency)\n",
    "\n",
    "H_graph_max = np.max(H_problem_diag)\n",
    "H_graph_min = np.min(H_problem_diag)\n",
    "\n",
    "print(H_problem_diag)\n",
    "print('H_max:', H_graph_max, '  H_min:', H_graph_min)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 搭建 QAOA 电路\n",
    "\n",
    "通过交替地摆放两个参数可调的电路模块，我们得以搭建QAOA电路\n",
    "\n",
    "$$U_x(\\beta_P)U_c(\\gamma_P)\\dots U_x(\\beta_1)U_c(\\gamma_1),$$\n",
    "\n",
    "其中放置的次数记为 $P$。具体地，模块一是由描述问题哈密顿量的矩阵生成，即\n",
    "\n",
    "$$U_c(\\gamma)=e^{-i \\gamma H_c },$$\n",
    "\n",
    "其中 $i= \\sqrt{-1}$ 是虚数单位， $\\gamma\\in [0, \\pi]$ 是可以调节的参数。模块二是\n",
    "\n",
    "$$U_x(\\beta)=e^{-i \\beta H_x },$$\n",
    "\n",
    "由描述驱动哈密顿量的另一个矩阵生成 \n",
    "\n",
    "$$H_x =X_1+X_2+X_3+X_4. $$\n",
    "\n",
    "$\\beta\\in [0, \\pi]$ 也是一个可调参数，算符 $X=\\begin{bmatrix} 0 & 1\\\\ 1 & 0\\end{bmatrix}$ 是作用在量子比特上的 Pauli-X 逻辑门，例如 $X_1$ 实际数学表达式为 $X_1\\otimes I_2\\otimes I_3\\otimes I_4$。\n",
    "\n",
    "QAOA 电路的每一模块可以进一步分解为若干个作用在单比特和两比特上的含参的量子逻辑门，如图所示:\n",
    "\n",
    "\n",
    "![QAOA.png](https://release-data.cdn.bcebos.com/PIC%2FQAOACir.png) \n",
    "\n",
    "\n",
    "其中，模块 $U_x(\\beta)$ 可以分解为在每个量子比特上作用绕 $X$ 方向转动的量子逻辑门 $R_x(\\beta)= e^{-i\\beta X_j}$，而模块 $U_c(\\gamma)$ 则可分解为作用在两比特上的 $ZZ$ 逻辑门 $R_{zz}(\\gamma)= e^{-i\\gamma Z\\otimes Z}$。\n",
    "\n",
    "此外，我们可以设置交叉放置两个模块的次数，记为 QAOA 电路的层数 $P$。于是输入\n",
    "\n",
    "- 量子电路的初始状态\n",
    "- 经典问题的邻接矩阵\n",
    "- 电路比特数目\n",
    "- 电路层数\n",
    "\n",
    "构建标准的 QAOA 量子电路："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def circuit_QAOA(theta, input_state, adjacency_matrix, N, P):\n",
    "    \"\"\"\n",
    "    This function constructs the parameterized QAOA circuit which is composed of P layers of two blocks:\n",
    "    one block is U_theta[layer][0] based on the problem Hamiltonian H which encodes the classical problem,\n",
    "    and the other is U_theta[layer][1] constructed from the driving Hamiltonian describing the rotation around Pauli X\n",
    "    acting on each qubit. It finally outputs the final state of the QAOA circuit.\n",
    "    \n",
    "    Args:\n",
    "         theta: parameters to be optimized in the QAOA circuit\n",
    "         input_state: initial state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings \n",
    "         in the computataional basis\n",
    "         adjacency_matrix:  the adjacency matrix of the graph encoding the classical problem\n",
    "         N: number of qubits, or equivalently, the number of parameters in the original classical problem\n",
    "         P: number of layers of two blocks in the QAOA circuit\n",
    "    Returns:\n",
    "        the final state of the QAOA circuit: cir.state\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    cir = UAnsatz(N, input_state=input_state)\n",
    "    \n",
    "    # This loop defines the QAOA circuit with P layers of two blocks\n",
    "    for layer in range(P):\n",
    "        # The second and third loops construct the first block U_theta[layer][0] which involves two-qubit operation\n",
    "        #  e^{-i\\beta Z_iZ_j} acting on a pair of qubits or nodes i and j in the circuit in each layer.\n",
    "        for row in range(N):\n",
    "            for col in range(N):\n",
    "                if abs(adjacency_matrix[row, col]) and row < col:\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "                    cir.rz(\n",
    "                        theta=theta[layer][0]\n",
    "                        * adjacency_matrix[row, col],\n",
    "                        which_qubit=col + 1,\n",
    "                    )\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "        # This loop constructs the second block U_theta only involving the single-qubit operation e^{-i\\beta X}.\n",
    "        for i in range(1, N + 1):\n",
    "            cir.rx(theta=theta[layer][1], which_qubit=i)\n",
    "\n",
    "    return cir.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在标准 QAOA 的基础上，我们还支持对电路结构进行扩展，进一步探索 QAOA 更多可能性。例如，可以将模块二的驱动哈密顿量 $H_x$ 中的的绕单比特 X 方向转动 $R_x(\\beta)$ 扩展为绕任意方向转动，且任意方向等价于依次绕 Z, X, Z 方向转动适当的角度，即$R_z(\\beta_1)R_x(\\beta_2)R_z(\\beta_3)$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def circuit_extend_QAOA(theta, input_state, adjacency_matrix, N, P):\n",
    "    \"\"\"\n",
    "    This is an extended version of the QAOA circuit, and the main difference is U_theta[layer]([1]-[3]) constructed\n",
    "    from the driving Hamiltonian describing the rotation around an arbitrary direction on each qubit.\n",
    "\n",
    "    Args:\n",
    "        theta: parameters to be optimized in the QAOA circuit\n",
    "        input_state: input state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings\n",
    "                     in the computational basis\n",
    "        adjacency_matrix:  the adjacency matrix of the problem graph encoding the original problem\n",
    "        N: number of qubits, or equivalently, the number of parameters in the original classical problem\n",
    "        P: number of layers of two blocks in the QAOA circuit\n",
    "    Returns:\n",
    "        final state of the QAOA circuit: cir.state\n",
    "\n",
    "    Note: If this U_extend_theta function is used to construct QAOA circuit, then we need to change the parameter layer\n",
    "           in the Net function defined below from the Net(shape=[D, 2]) for U_theta function to Net(shape=[D, 4])\n",
    "           because the number of parameters doubles in each layer in this QAOA circuit.\n",
    "    \"\"\"\n",
    "    cir = UAnsatz(N, input_state=input_state)\n",
    "\n",
    "    for layer in range(P):\n",
    "        for row in range(N):\n",
    "            for col in range(N):\n",
    "                if abs(adjacency_matrix[row, col]) and row < col:\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "                    cir.rz(\n",
    "                        theta=theta[layer][0]\n",
    "                        * adjacency_matrix[row, col],\n",
    "                        which_qubit=col + 1,\n",
    "                    )\n",
    "                    cir.cnot([row + 1, col + 1])\n",
    "\n",
    "        for i in range(1, N + 1):\n",
    "            cir.rz(theta=theta[layer][1], which_qubit=i)\n",
    "            cir.rx(theta=theta[layer][2], which_qubit=i)\n",
    "            cir.rz(theta=theta[layer][3], which_qubit=i)\n",
    "\n",
    "    return cir.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "搭建 QAOA 量子电路的工作完成后，此时量子电路的输出状态为\n",
    "\n",
    "$$|\\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)\\rangle=U_x(\\beta_P)U_c(\\gamma_P)\\dots U_x(\\beta_1)U_c(\\gamma_1)|+\\rangle_1\\dots|+\\rangle_N.$$\n",
    "\n",
    "其中每个量子比特的初始状态处于量子叠加态 $|+\\rangle=\\frac{1}{\\sqrt{2}}\\left(|0\\rangle+|1\\rangle\\right)$ 。最终，我们得到量子优化问题的损失函数\n",
    "\n",
    "$$F_P=\\min_{\\boldsymbol{\\beta},\\boldsymbol{\\gamma}} \\langle \\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)| H_c|\\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)\\rangle.$$\n",
    "\n",
    "因为 QAOA 是一个量子-经典混杂算法，所以搭建完成 QAOA 电路且得到相应的损失函数后，我们可以进一步利用经典的优化算法寻找最优参数 $\\boldsymbol{\\beta},\\boldsymbol{\\gamma}$，从而形成一个完整的闭环网络。\n",
    "\n",
    "下面的函数给出了通过 Paddle Quantum 搭建的完整 QAOA 网络："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class Net(fluid.dygraph.Layer):\n",
    "    \"\"\"\n",
    "    It constructs the net for QAOA which combines the  QAOA circuit with the classical optimizer which sets rules\n",
    "    to update parameters described by theta introduced in the QAOA circuit.\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        shape,\n",
    "        param_attr=fluid.initializer.Uniform(low=0.0, high=np.pi, seed=1024),\n",
    "        dtype=\"float32\",\n",
    "    ):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.theta = self.create_parameter(\n",
    "            shape=shape, attr=param_attr, dtype=dtype, is_bias=False\n",
    "        )\n",
    "\n",
    "    def forward(self, input_state, adjacency_matrix, out_state_store, N, P, METHOD):\n",
    "        \"\"\"\n",
    "        This function constructs the loss function for the QAOA circuit.\n",
    "\n",
    "        Args:\n",
    "            self: the free parameters to be optimized in the QAOA circuit and defined in the above function\n",
    "            input_state: initial state of the QAOA circuit which usually is the uniform superposition of 2^N bit-strings\n",
    "                         in the computational basis $|0\\rangle, |1\\rangle$\n",
    "            adjacency_matrix: the adjacency matrix generated from the graph encoding the classical problem\n",
    "            out_state_store: the output state of the QAOA circuit\n",
    "            N: number of qubits\n",
    "            P: number of layers\n",
    "            METHOD: which version of QAOA is chosen to solve the problem, i.e., standard version labeled by 1 or\n",
    "            extended version by 2.\n",
    "        Returns:\n",
    "            The loss function for the parameterized QAOA circuit.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Generate the problem_based quantum Hamiltonian H_problem based on the classical problem in paddle\n",
    "        H, _ = H_generator(N, adjacency_matrix)\n",
    "        H_problem = fluid.dygraph.to_variable(H)\n",
    "\n",
    "        # The standard QAOA circuit: the function circuit_QAOA is used to construct the circuit, indexed by METHOD 1.\n",
    "        if METHOD == 1:\n",
    "            out_state = circuit_QAOA(self.theta, input_state, adjacency_matrix, N, P)\n",
    "        # The extended QAOA circuit: the function circuit_extend_QAOA is used to construct the net, indexed by METHOD 2.\n",
    "        elif METHOD == 2:\n",
    "            out_state = circuit_extend_QAOA(self.theta, input_state, adjacency_matrix, N, P)\n",
    "        else:\n",
    "            raise ValueError(\"Wrong method called!\")\n",
    "\n",
    "        out_state_store.append(out_state.numpy())\n",
    "        loss = pp_matmul(\n",
    "            pp_matmul(out_state, H_problem),\n",
    "            transpose(\n",
    "                fluid.framework.ComplexVariable(out_state.real, -out_state.imag),\n",
    "                perm=[1, 0],\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        return loss.real\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 训练网络\n",
    "\n",
    "我们开始训练整个 QAOA 网络，即通过优化参数向量 $\\boldsymbol{\\beta}=(\\beta_1,\\beta_2,\\beta_3,\\beta_4)$ 和 $\\boldsymbol{\\gamma}=(\\gamma_1,\\gamma_2,\\gamma_3, \\gamma_4)$ 来达到求解 $H_c$ 最小特征值的目的。\n",
    "\n",
    "与经典机器学习算法一样，首先设置 QAOA 网络里的超参数:\n",
    "\n",
    "- 电路比特数目 N\n",
    "- 电路的层数 P\n",
    "- 迭代次数 ITR\n",
    "- 学习步长 LR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "N = 4 # number of qubits, or number of nodes in the graph\n",
    "P = 4 # number of layers \n",
    "ITR = 120  # number of iteration steps\n",
    "LR = 0.1  # learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，灵活调用：\n",
    "\n",
    "- 量子电路初始状态：每个量子比特态处于相干叠加态 $\\frac{1}{\\sqrt{2}}\\left(|0\\rangle+|1\\rangle\\right)$\n",
    "- 采用标准 QAOA （记为 METHOD=1）或者扩展 QAOA （记为 METHOD = 2）\n",
    "- 经典优化器 Adam optimizer \n",
    "\n",
    "最后，训练模型并保存结果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def Paddle_QAOA(classical_graph_adjacency, N, P, METHOD, ITR, LR):\n",
    "    \"\"\"\n",
    "    This is the core function to run QAOA.\n",
    "\n",
    "     Args:\n",
    "         classical_graph_adjacency: adjacency matrix to describe the graph which encodes the classical problem\n",
    "         N: number of qubits (default value N=4)\n",
    "         P: number of layers of blocks in the QAOA circuit (default value P=4)\n",
    "         METHOD: which version of the QAOA circuit is used: 1, standard circuit (default); 2, extended circuit\n",
    "         ITR: number of iteration steps for QAOA (default value ITR=120)\n",
    "         LR: learning rate for the gradient-based optimization method (default value LR=0.1)\n",
    "     Returns:\n",
    "         optimized parameters theta and the bitstrings sampled from the output state with maximal probability\n",
    "    \"\"\"\n",
    "    out_state_store = []\n",
    "    with fluid.dygraph.guard():\n",
    "        # Preparing the initial state\n",
    "        _initial_state = np.ones([1, 2 ** N]).astype(\"complex64\") / np.sqrt(2 ** N)\n",
    "        initial_state = fluid.dygraph.to_variable(_initial_state)\n",
    "\n",
    "        # Construct the net or QAOA circuits based on the standard modules\n",
    "        if METHOD == 1:\n",
    "            net = Net(shape=[P, 2])\n",
    "        # Construct the net or QAOA circuits based on the extended modules\n",
    "        elif METHOD == 2:\n",
    "            net = Net(shape=[P, 4])\n",
    "        else:\n",
    "            raise ValueError(\"Wrong method called!\")\n",
    "\n",
    "        # Classical optimizer\n",
    "        opt = fluid.optimizer.AdamOptimizer(learning_rate=LR, parameter_list=net.parameters())\n",
    "\n",
    "        # Gradient descent loop\n",
    "        summary_iter, summary_loss = [], []\n",
    "        for itr in range(1, ITR + 1):\n",
    "            loss = net(\n",
    "                initial_state, classical_graph_adjacency, out_state_store, N, P, METHOD\n",
    "            )\n",
    "            loss.backward()\n",
    "            opt.minimize(loss)\n",
    "            net.clear_gradients()\n",
    "\n",
    "            print(\"iter:\", itr, \"  loss:\", \"%.4f\" % loss.numpy())\n",
    "            summary_loss.append(loss[0][0].numpy())\n",
    "            summary_iter.append(itr)\n",
    "\n",
    "        theta_opt = net.parameters()[0].numpy()\n",
    "        print(theta_opt)\n",
    "\n",
    "        os.makedirs(\"output\", exist_ok=True)\n",
    "        np.savez(\"./output/summary_data\", iter=summary_iter, energy=summary_loss)\n",
    "\n",
    "    # Output the measurement probability distribution which is sampled from the output state of optimized QAOA circuit.\n",
    "    prob_measure = np.zeros([1, 2 ** N]).astype(\"complex\")\n",
    "\n",
    "    rho_out = out_state_store[-1]\n",
    "    rho_out = np_matmul(np.conjugate(rho_out).T, rho_out).astype(\"complex\")\n",
    "\n",
    "    for index in range(0, 2 ** N):\n",
    "        comput_basis = np.zeros([1, 2 ** N])\n",
    "        comput_basis[0][index] = 1\n",
    "        prob_measure[0][index] = np.real(np_matmul(np_matmul(comput_basis, rho_out), comput_basis.T))\n",
    "\n",
    "    return prob_measure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调用模型训练结果，输出得到的最优参数向量 $\\boldsymbol{\\beta}^*$ 和 $\\boldsymbol{\\gamma}^*$，并且将 QAOA 的输出结果和真实结果进行比较："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "classical_graph, classical_graph_adjacency = generate_graph(N, 1)\n",
    "\n",
    "prob_measure_dis = Paddle_QAOA(classical_graph_adjacency, N =4, P=4, METHOD=1, ITR=120, LR=0.1)\n",
    "\n",
    "# Load the data of QAOA\n",
    "x1 = np.load('./output/summary_data.npz')\n",
    "\n",
    "H_min = np.ones([len(x1['iter'])]) * H_graph_min\n",
    "\n",
    "# Plot it\n",
    "\n",
    "loss_QAOA, = plt.plot(x1['iter'], x1['energy'], \\\n",
    "                                        alpha=0.7, marker='', linestyle=\"--\", linewidth=2, color='m')\n",
    "benchmark, = plt.plot(x1['iter'], H_min, alpha=0.7, marker='', linestyle=\":\", linewidth=2, color='b')\n",
    "plt.xlabel('Number of iteration')\n",
    "plt.ylabel('Performance of the loss function for QAOA')\n",
    "\n",
    "plt.legend(handles=[\n",
    "    loss_QAOA,\n",
    "    benchmark\n",
    "],\n",
    "    labels=[\n",
    "            r'Loss function $\\left\\langle {\\psi \\left( {\\bf{\\theta }} \\right)} '\n",
    "            r'\\right|H\\left| {\\psi \\left( {\\bf{\\theta }} \\right)} \\right\\rangle $',\n",
    "            'The benchmark result',\n",
    "    ], loc='best')\n",
    "\n",
    "# Show the picture\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## 5. 解码量子答案\n",
    "\n",
    "当求得损失函数 $\\langle \\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)| H_{\\rm Cut}|\\psi(\\boldsymbol{\\beta},\\boldsymbol{\\gamma}, P)\\rangle$ 的最小值以及相对应的一组参数 $(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*)$ 后，我们的任务还没有完成。为了进一步求得 Max-Cut 问题的解，需要从 QAOA 输出的量子态 \n",
    "\n",
    "$$|\\psi(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*, P)\\rangle=\\sum_{i=1}^{2^4}\\lambda_i |\\boldsymbol{x}_i\\rangle$$\n",
    "\n",
    "中解码出经典优化问题的答案。上式中  $\\boldsymbol{x}_i=x_1x_2x_3 x_4\\in \\{0, 1\\}^4$，对应着经典问题的一个可行解。物理上，解码量子态需要对量子态进行测量，然后统计测量结果的概率分布：\n",
    "                 \n",
    "$$ p(\\boldsymbol{x})=|\\langle \\boldsymbol{x}|\\psi(\\boldsymbol{\\beta}^*,\\boldsymbol{\\gamma}^*,P)\\rangle|^2.$$\n",
    "              \n",
    "\n",
    "某种程度上，某个比特串出现的概率越大，意味着其对应的 Max-Cut 问题最优解的可能性越大。\n",
    "\n",
    "此外，Paddle Quantum 提供了查看 QAOA 量子电路输出状态的测量结果概率分布的函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prob_measure = prob_measure_dis.flatten()\n",
    "\n",
    "pos = nx.circular_layout(classical_graph)\n",
    "# when N is large, it is not suggested to plot this figure\n",
    "name_list = [np.binary_repr(index, width=N) for index in range(0, 2 ** N)]\n",
    "plt.bar(\n",
    "        range(len(np.real(prob_measure))),\n",
    "        np.real(prob_measure),\n",
    "        width=0.7,\n",
    "        tick_label=name_list,\n",
    "    )\n",
    "plt.xticks(rotation=90)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，再次利用参数代换 $|x \\rangle\\rightarrow z=2x-1\\in\\{-1, 1\\}$，可以从量子答案中解码得到 Max-Cut 问题的可行解。 此时，记 $z_i=-1$ 的顶点属于集合 $S^\\prime$ 以及 $z_j=1$ 的顶点属于集合 $S$，这两个顶点集合之间存在的边就是该图的一个可能得最大割方案。 \n",
    "\n",
    "选取测量结果中出现几率最大的比特串，然后将其映射回经典解，并且画出对应的最大割方案：\n",
    "\n",
    "- 蓝色顶点属于集合 $S$\n",
    "- 红色顶点属于集合 $S^\\prime$\n",
    "- 折线属于图的一条割线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# Find the position of max value in the measure_prob_distribution\n",
    "max_prob_pos_list = np.where(prob_measure == max(prob_measure))\n",
    "# Store the max value from ndarray to list\n",
    "max_prob_list = max_prob_pos_list[0].tolist()\n",
    "# Change it to the  binary format\n",
    "solution_list = [np.binary_repr(index, width=N) for index in max_prob_list]\n",
    "print(\"The output bitstring:\", solution_list)\n",
    "\n",
    "# Draw the graph representing the first bitstring in the solution_list to the MaxCut-like problem\n",
    "head_bitstring = solution_list[0]\n",
    "\n",
    "node_cut = [\"blue\" if head_bitstring[node] == \"1\" else \"red\" for node in classical_graph]\n",
    "\n",
    "edge_cut = [\n",
    "    \"solid\" if head_bitstring[node_row] == head_bitstring[node_col] else \"dashed\"\n",
    "    for node_row, node_col in classical_graph.edges()\n",
    "    ]\n",
    "nx.draw(\n",
    "        classical_graph,\n",
    "        pos,\n",
    "        node_color=node_cut,\n",
    "        style=edge_cut,\n",
    "        width=4,\n",
    "        with_labels=True,\n",
    "        font_weight=\"bold\",\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "source": [
    "# 参考文献\n",
    "\n",
    "[1] E. Farhi, J. Goldstone, and S. Gutman. 2014. A quantum approximate optimization algorithm. arXiv:1411.4028 "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
